﻿#ifndef __JS_MODEL_PROXY_H__
#define __JS_MODEL_PROXY_H__
#include "base64.h"
#include "models/SoftMaxModel.h"
#include "RoadSignModel.h"
#include "utils/utils.h"
#include "zip.h"
#include <map>

#ifdef __EMSCRIPTEN__
#include <emscripten.h>
#include <emscripten/bind.h>
#endif

using SoftMaxModel = darknet::SoftMaxModel;
using DataImage    = darknet::DataImage;
using namespace darknet;
class JsModelProxy
{
public:
    JsModelProxy(SoftMaxModel *model) : m_model(model) {}

    std::vector<std::string> GetClassifyList()
    {
        return std::move(GetModel()->GetClassifyList());
    }

    std::string AddTrainData(const std::string &path, const std::string &classify, const std::string &base64_image_data)
    {
        auto img = GetImageByBase64ImageData(base64_image_data);
        if (img == nullptr) return "";
        auto fileppath = std::filesystem::path(path).replace_extension(".jpg").string();
        img->Save(fileppath);
        GetModel()->SetPathClassifyItem(fileppath, classify);
        GetModel()->SaveDataSet();
        FlushFileSystem();

        return fileppath;
    }

    void SetTrainData(const std::string &path, const std::string &classify)
    {
        GetModel()->SetPathClassifyItem(path, classify);
        GetModel()->SaveDataSet();
        FlushFileSystem();
    }

    void DeleteTrainData(const std::string &path)
    {
        GetModel()->UnsetPathClassifyItem(path);
        DeleteFile(path);
        GetModel()->SaveDataSet();
        FlushFileSystem();
    }

    std::map<std::string, std::string> LoadTrainData()
    {
        GetModel()->LoadDataSet();

        std::map<std::string, std::string> datas;
        auto dataset = GetModel()->GetDataSet();
        for (auto it = dataset->begin(); it != dataset->end(); it++)
        {
            datas.insert_or_assign(it->first, it->second->GetClassify());
        }

        // std::cout << "LoadTrainData => " << GetModel()->GetDataSetPath() << " size = " << datas.size() << std::endl;
        // std::cout << "exist = " << std::filesystem::exists(GetModel()->GetDataSetPath()) << std::endl;
        return std::move(datas);
    }

    void SaveTrainData()
    {
        GetModel()->SaveDataSet();
        FlushFileSystem();
    }

    void LoadModel(const std::string &path)
    {
        GetModel()->LoadModelFromFile(path);
    }

    void SaveModel()
    {
        GetModel()->SaveModelToFile(GetModel()->GetModelPath());
        FlushFileSystem();
    }

    void DeleteModel()
    {
        DeleteFile(GetModel()->GetModelPath());
        FlushFileSystem();
    }

    std::string GetModelName() { return GetModel()->GetModelName(); }

    void SetRootDirectory(const std::string &root_directory)
    {
        auto old_root_directory = GetModel()->GetRootDirectory();
        // auto old_dataset_directory = GetModel()->GetDataSetDirectory();
        // auto old_dataset_filepath = GetModel()->GetDataSetPath();
        // auto old_image_directory = GetImageDirectory();

        GetModel()->SetRootDirectory(root_directory);
        GetModel()->SetModelDirectory((std::filesystem::path(root_directory) / "models").string());
        GetModel()->SetDataSetDirectory((std::filesystem::path(root_directory) / "datasets").string());
        GetModel()->SetModelPath((std::filesystem::path(GetModel()->GetModelDirectory()) / GetModel()->GetDefaultModelFileName()).string());
        GetModel()->SetDataSetPath((std::filesystem::path(GetModel()->GetDataSetDirectory()) / ("dataset.txt")).string());

        std::filesystem::create_directories(GetModel()->GetModelDirectory());
        std::filesystem::create_directories(GetModel()->GetDataSetDirectory());
        std::filesystem::create_directories(GetImageDirectory());

        std::string assets_root_directory  = "/assets/";
        std::string darkent_root_directory = "/darknet/";
        if (old_root_directory.find_first_of(assets_root_directory) == 0)
        {
            GetModel()->LoadDataSet(); // 加载数据集
            auto datasets = GetModel()->GetDataSet();
            for (auto it = datasets->begin(); it != datasets->end(); it++)
            {
                auto path = it->first;
                if (path.find_first_of(darkent_root_directory) == 0 && !std::filesystem::exists(path))
                {
                    auto src_path = assets_root_directory + path.substr(darkent_root_directory.size());
                    std::cout << "copy file: " << src_path << " => " << path << std::endl;
                    if (std::filesystem::exists(src_path)) std::filesystem::copy(src_path, path);
                }
            }
            GetModel()->SaveDataSet(); // 保存合并后数据集
        }

        // std::cout << "root directory: " << root_directory << std::endl;
        // std::cout << "model directory: " << GetModel()->GetModelDirectory() << std::endl;
        // std::cout << "data set directory: " << GetModel()->GetDataSetDirectory() << std::endl;
        // std::cout << "image directory: " << GetImageDirectory() << std::endl;
        // std::cout << "model filepath: " << GetModel()->GetModelPath() << std::endl;
        // std::cout << "dataset filepath: " << GetModel()->GetDataSetPath() << std::endl;
        if (std::filesystem::exists(GetModel()->GetModelPath()))
        {
            if (GetModel()->LoadModelFromFile(GetModel()->GetModelPath()))
            {
                std::cout << "load custom model form " << GetModel()->GetModelPath() << " success" << std::endl;
            }
        }
        FlushFileSystem();
    }

    void SetTrainParams(int batch, bool random)
    {
        GetModel()->SetBatch(batch);
        GetModel()->GetNet()->SetBatch(batch);
        GetModel()->SetRandom(random);
        GetModel()->LoadDataSet();
        GetModel()->RefreshNet();
        GetModel()->PrintModelInfo();
    }

    float DoEpoch()
    {
        return GetModel()->DoEpoch();
    }

    std::string Predict(const std::string &base64_image_data)
    {
        auto start_time   = get_timestamp();
        float probability = 0.0f;
        auto img          = GetImageByBase64ImageData(base64_image_data);
        auto index        = GetModel()->Predict(img, &probability);
        auto classify     = GetModel()->GetIndexClassify(index);
        std::ostringstream oss;
        oss << classify << " " << probability;
        if (GetModel() == RoadSignModel::GetInstance())
        {
            auto roadsign = RoadSignModel::GetInstance();
            oss << " " << roadsign->GetX() << " " << roadsign->GetY() << " " << roadsign->GetWidth() << " " << roadsign->GetHeight();
        }
        auto end_time = get_timestamp();
        std::cout << GetModel()->GetModelName() <<  " predict use time: " << (end_time - start_time) << "ms" << ", classify: " << classify << ", probability: " << probability << std::endl;
        return oss.str();
    }

    std::string DownloadDataSet()
    {
        auto zip_path       = GetDataSetZipPath();
        auto root_directory = GetModel()->GetRootDirectory();
        return PackDirectory(zip_path, root_directory);
    }

#ifdef __EMSCRIPTEN__
    emscripten::val JsDownloadDataSet() { return StringToBuffer(std::move(DownloadDataSet())); }
#endif

public:
    inline SoftMaxModel *GetModel() { return m_model; }
    inline std::string GetDataSetZipPath() { return GetModel()->GetRootDirectory() + "/dataset.zip"; }
    inline std::string GetImageDirectory() { return GetModel()->GetDataSetDirectory() + "/images"; }

protected:
    SoftMaxModel *m_model;

public:
    static std::shared_ptr<DataImage> GetImageByImagePath(const std::string &image_path)
    {
        auto img = std::make_shared<DataImage>();
        if (!img->Load(image_path))
        {
            std::cout << "invalid image path!!!" << image_path << std::endl;
            return nullptr;
        }
        return img;
    }

    static unsigned char* LoadImageRGBByBase64ImageData(const std::string &base64_image_data, int* width, int* height)
    {
        std::string jpeg_prefix = "data:image/jpeg;base64,";
        std::string png_prefix  = "data:image/png;base64,";
        std::string img_data;
        if (base64_image_data.substr(0, jpeg_prefix.length()) == jpeg_prefix)
        {
            img_data = base64_decode(base64_image_data.substr(jpeg_prefix.length()));
        }
        else if (base64_image_data.substr(0, png_prefix.length()) == png_prefix)
        {
            img_data = base64_decode(base64_image_data.substr(png_prefix.length()));
        }
        else
        {
            img_data = base64_decode(base64_image_data);
        }
        return DataImage::LoadRGB((unsigned char *)(img_data.c_str()), img_data.size(), width, height);
    }
    
    static std::shared_ptr<DataImage> GetImageByBase64ImageData(const std::string &base64_image_data)
    {
        std::string jpeg_prefix = "data:image/jpeg;base64,";
        std::string png_prefix  = "data:image/png;base64,";
        std::string img_data;
        if (base64_image_data.substr(0, jpeg_prefix.length()) == jpeg_prefix)
        {
            img_data = base64_decode(base64_image_data.substr(jpeg_prefix.length()));
        }
        else if (base64_image_data.substr(0, png_prefix.length()) == png_prefix)
        {
            img_data = base64_decode(base64_image_data.substr(png_prefix.length()));
        }
        else
        {
            img_data = base64_decode(base64_image_data);
        }
        auto img = std::make_shared<DataImage>();
        if (!img->Load((unsigned char *)(img_data.c_str()), img_data.size()))
        {
            std::cout << "invalid image data!!!" << std::endl;
            return nullptr;
        }
        return img;
    }

    static bool SaveFile(const std::string &path, const std::string &data)
    {
        std::ofstream ofs(path, std::ios::binary);
        if (ofs)
        {
            ofs << data;
        }
        else
        {
            std::cout << "invalid path: " << path << std::endl;
            return false;
        }
        FlushFileSystem();
        return true;
    }

    static std::string LoadFile(const std::string &path)
    {
        std::ifstream ifs(path, std::ios::binary);
        if (ifs)
        {
            std::string text((std::istreambuf_iterator<char>(ifs)), (std::istreambuf_iterator<char>()));
            ifs.close();
            return text;
        }
        else
        {
            std::cout << "invalid path: " << path << std::endl;
            return "";
        }
    }

    static void DeleteFile(const std::string &path)
    {
        std::filesystem::remove(path);
        FlushFileSystem();
    }

    static void FlushFileSystem()
    {
#ifdef __EMSCRIPTEN__
        EM_ASM({
            setTimeout(function() { 
                if (window.FS_syncfs_running) return;
                window.FS_syncfs_running = true;
                FS.syncfs(false, function() {
                    window.FS_syncfs_running = false;
                }); }, 100);
        });
#endif
    }

#ifdef __EMSCRIPTEN__
    static emscripten::val StringToBuffer(const std::string &str)
    {
        const uint8_t *data    = reinterpret_cast<const uint8_t *>(str.data());
        size_t size            = str.size();
        emscripten::val buffer = emscripten::val(emscripten::typed_memory_view(size, data));
        return buffer;
    }
#endif

    // 添加单个文件到 ZIP 存档
    static bool AddFileToZip(zip_t *archive, const std::filesystem::path &filepath, const std::string &prefix)
    {
        zip_source_t *source = zip_source_file(archive, filepath.string().c_str(), 0, 0);
        if (source == nullptr)
        {
            std::cerr << "Failed to open file: " << filepath << std::endl;
            return false;
        }
        if (zip_file_add(archive, (prefix + filepath.filename().string()).c_str(), source, ZIP_FL_ENC_UTF_8) < 0)
        {
            std::cerr << "Failed to add file to zip archive." << std::endl;
            zip_source_free(source);
            return false;
        }
        return true;
    }

    // 递归添加目录到 ZIP 存档
    static bool AddDirectoryToZip(zip_t *archive, const std::filesystem::path &directory, const std::string &prefix)
    {
        for (const auto &entry : std::filesystem::directory_iterator(directory))
        {
            if (entry.is_directory())
            {
                if (!AddDirectoryToZip(archive, entry.path(), prefix + entry.path().filename().string() + "/"))
                {
                    return false;
                }
            }
            else
            {
                if (!AddFileToZip(archive, entry.path(), prefix))
                {
                    return false;
                }
            }
        }
        return true;
    }

    // 打包目录
    static std::string PackDirectory(const std::string &out_filepath, const std::string &in_directory)
    {
        auto archive = zip_open(out_filepath.c_str(), ZIP_CREATE | ZIP_TRUNCATE, 0);
        if (archive == nullptr) return "";

        if (!AddDirectoryToZip(archive, in_directory, ""))
        {
            std::cout << "添加文件到zip文件时出错:" << zip_strerror(archive) << std::endl;
            return "";
        }

        if (-1 == zip_close(archive))
        {
            std::cout << "关闭zip文件时出错:" << zip_strerror(archive) << std::endl;
            return "";
        }

        auto zip_data = LoadFile(out_filepath);
        std::cout << "DownloadTrainerData Success, Size = " << zip_data.size() << std::endl;
        std::filesystem::remove(out_filepath);
        return std::move(zip_data);
    }
};

#endif